import numpy as np
import os
import torch
import pywt

from ENF_deepLearning.utils.helper.Data_Helper import Data_Helper
from ENF_deepLearning.utils.resnet.ResNet_train import ResNet_train

import d2lzh_pytorch as d2l
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def dwt(data, count=1, db='db4'):
    ca = data
    cd = 0
    for i in range(count):
        ca, cd = pywt.dwt(ca, db)
    return ca, cd

data_helper = Data_Helper()  # 数据相关的操作助手
resnet = ResNet_train()  # cnn分类器
data_path = '../utils/data'
list = os.listdir(data_path)

Y = []
X = []
minn = 0
PAD = 350
for item in list:
    path = data_path + '/' + item
    # print(item)
    data = np.loadtxt(path, dtype=np.float, delimiter=',')
    minn = max(minn, data.shape[0])
    data = data.tolist()
    if len(data) < PAD:
        for i in range(PAD - len(data)):
            data.append(0)
    # plt.plot(data)
    # plt.show()
    X.append(np.array(data[1:]).astype(np.float))  # x表示值
    Y.append(np.array(data[0]).astype(np.int))   # y表示标签

X = np.array(X)
Y = np.array(Y)

x_train, x_test, y_train, y_test = data_helper.train_test_split(X, Y, 0.1, True)   ## 划分为仅有10%的测试集

print(x_train.shape, x_test.shape, y_train.shape, y_test.shape)



avg_acc = resnet.one_D_resNet(x_train.reshape(-1, 1, PAD-1), x_test, y_train, y_test, epoch=1500, batch_size=32, lr=0.0001, print_for_epoch=10,
                        n_filter=64, kernel_size=128, stride=1, number_class=2, filter_row=1,#65, 60
                        loss_function=torch.nn.CrossEntropyLoss, optimizer=torch.optim.SGD)    # 1D-resnet训练

